from database.database import text_embedding
import numpy as np
import pickle


class rag:
    def __init__(self, question, database_path = './database/data.pkl'):
        '''利用余弦相似度召回知识库语料，同时以相似度作为排序依据，返回前三个得分最高的知识构建提示词'''
        self.question = question
        with open(database_path, 'rb') as f:
            self.v_database = pickle.load(f)
        self.set_prompt()
        self.search_examples = []

    def binary_search(self, arr, target):
        left, right = 0, len(arr) - 1
        while left <= right:
            mid = (left + right) // 2
            if arr[mid]['id'] == target:
                return mid
            elif arr[mid]['id'] > target:
                right = mid - 1
            else:
                left = mid + 1
        return -1

    def set_score(self):
        a = text_embedding(self.question, bert_path = './bert-base-chinese')
        text_score = []
        for text in self.v_database:
            b = text['text_embedding']
            similarity = np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
            text_score.append((text['id'], similarity))
        text_score.sort(key = lambda x: x[-1], reverse = True)
        return text_score

    def set_prompt(self):
        score_list = self.set_score()[:3]
        for id, score in score_list:
            result = self.binary_search(self.v_database, id)
            if result != -1:
                self.search_examples.append(self.v_database[result]['text'])
        return self.search_examples
